# -*- coding: utf-8 -*-
"""
核心内容
包括数据格式，数据运算，自动微分
"""

import numpy as np
import weakref
import contextlib
import dezero

try:
    import cupy
    array_types = (np.ndarray, cupy.ndarray)
except ImportError:
    array_types = (np.ndarray)


#=========================Config===========================

class Config:
    """单例，控制禁用/启用反向传播模式"""
    enable_backprop = True
    train = True
    
#为了使用with语句使用
@contextlib.contextmanager
def using_config(name, value):
    old_Value = getattr(Config, name)
    setattr(Config, name, value)
    try:
        yield
    finally:
        setattr(Config, name, old_Value)

def test_mode():
    """进入测试模式时暂时关闭训练模式"""
    return using_config('train', False)
        
def no_grad():
    """配合with使用，在with中的程序可以禁用反向传播模式，退出后自动启用"""
    return using_config('enable_backprop',False)

#===========================Variable==============================
#基本变量类
class Variable:
    """
    参数：data, name
    功能：基本变量类,data保存变量的值,grad保存变量的导数
    """
    
    __array_priority__ = 200 #将var的运算符优先级高于ndarray的运算符优先级
    
    def __init__(self, data, name=None):
        #data只支持ndarray格式
        if data is not None:
            if not isinstance(data, array_types): #输入支持np和cp类型
                raise TypeError('{} is not supported'.format(type(data)))
        
        #实际数据保存在data中
        self.data = data
        #添加name变量，为计算图可视化的标签
        self.name = name
        #保存导数值
        self.grad = None
        #保存这个变量的创建者
        self.creator = None
        #保存自己的优先级变量,自己在第几层
        self.generation = 0
    
    def set_creator(self,func):
        """设置这个变量的创建者"""
        self.creator = func
        self.generation = func.generation + 1
        
    def cleargrad(self):
        """重置导数"""
        self.grad = None
        
    def backward(self, retain_grad=False, create_graph=False):
        """
        输入：
            retain_grad：是否保存导数，不保存导数可优化速度，但是无法生成计算图
            create_graph：是否在反向传播时也创建计算图即求二阶导数，若要求二阶导数则设为True
        功能：自动连续求导，结束后可得到输入参数的导数
        """
        if self.grad is None:
            xp = dezero.cuda.get_array_module(self.data) #判断data类型
            self.grad = Variable(xp.ones_like(self.data)) #将导数也设为variable则可以进行高阶求导
        
        #按辈份取函数
        funcs = []
        seen_set = set()
        
        def add_func(f):
            if f not in seen_set:
                seen_set.add(f)
                funcs.append(f)
                funcs.sort(key=lambda x: x.generation)
                
        add_func(self.creator)
        
        while funcs:
            f = funcs.pop()
            gys = [output().grad for output in f.outputs] #注意这里的output是弱引用

            with using_config('enable_backprop', create_graph):  # 若create_graph为false则不在反向传播中进行反向传播即求二阶导数
                gxs = f.backward(*gys)
                if not isinstance(gxs,tuple):
                    gxs = (gxs,)

                #给前一个Variable的导数赋值
                for x,gx in zip(f.inputs, gxs):
                    #防止重复使用变量时,导数被覆盖,应该相加
                    if x.grad is None:
                        x.grad = gx
                    else:
                        x.grad = x.grad + gx

                    if x.creator is not None:
                        add_func(x.creator)
            
            #不保留个函数输出变量的导数
            if not retain_grad:
                for y in f.outputs:
                    y().grad = None #y是weakref
                    
    """
    使Variable使用上更像ndarray
    """
    @property
    def ndim(self):
        return self.data.ndim
    @property
    def size(self):
        return self.data.size
    @property
    def dtype(self):
        return self.data.dtype
    @property
    def shape(self):
        return self.data.shape

    def reshape(self, *shape):
        '''向reshape传入可变长参数，像np一样调用'''
        if len(shape) == 1 and isinstance(shape[0],(tuple,list)):
            shape = shape[0]
        return dezero.functions.reshape(self,shape)

    def transpose(self, *axes):
        if len(axes) == 0:
            axes = None
        elif len(axes) == 1:
            if isinstance(axes[0], (tuple, list)) or axes[0] is None:
                axes = axes[0]
        return dezero.functions.transpose(self, axes)
    @property
    def T(self):
        return dezero.functions.transpose(self)

    def sum(self,axis=None,keepdims=False):
        """求和函数，可以指定某个轴求和"""
        return dezero.functions.sum(self,axis,keepdims)

    def __len__(self):
        return len(self.data)
    def __repr__(self):
        if self.data is None:
            return 'variable(None)'
        p = str(self.data).replace('\n','\n'+' '*9)
        return 'variable(' + p + ')'

    def to_cpu(self):
        """cupy->numpy"""
        if self.data is not None:
            self.data = dezero.cuda.as_numpy(self.data)
    def to_gpu(self):
        """numpy->cupy"""
        if self.data is not None:
            self.data = dezero.cuda.as_cupy(self.data)

#检查是否ndarray,是标量则变成array输出
def as_array(x, array_module=np):
    if np.isscalar(x):
        return array_module.array(x)
    return x
#将ndarray转换成variable使用
def as_variable(obj):
    if isinstance(obj,Variable):
        return obj
    return Variable(obj)

    
#=================================Function======================
#方法基类
class Function:
    """
    方法基类
    子类实现forward 和 backward实现
    输入是 ndarray 列表 输出为元组
    inputs.data -> Function.forward -> outputs.data
    inputs.grad <- Function.backward <- outputs.grad
    """
    def __call__(self,*inputs):
        #将输入的参数如ndarray都变为variable使用
        inputs = [as_variable(x) for x in inputs]
        
        #input 为 Variable
        xs = [x.data for x in inputs]
        #具体的计算在forward方法中进行
        ys = self.forward(*xs)  #参数解包
        if not isinstance(ys,tuple): 
            ys = (ys,) #将输出变成元组
        outputs = [Variable(as_array(y)) for y in ys]
        
        #判断是否为反向传播模式
        if Config.enable_backprop:
            #func的辈份值是取输入的变量中最大的辈份值
            self.generation = max([x.generation for x in inputs])
            #让输出变量保存创造者信息
            for output in outputs:
                output.set_creator(self)
                
            self.inputs = inputs #保存输入的变量
            
            #这里outputs使用弱引用，因为反向传播后这个数就不需要了
            self.outputs = [weakref.ref(output) for output in outputs] #保存输出的变量
        
        #如果列表中只有1个元素,则返回第一个元素,如果函数返回值只有一个则变量直接返回
        return outputs if len(outputs)>1 else outputs[0]

    #具体的计算
    def forward(self,x):
        #如果子类不实现这个接口则报错
        raise NotImplementedError() 
    
    #计算导数的反向传播功能
    def backward(self,gy):
        #gy是从输出传播而来的导数
        raise NotImplementedError()
    
    
#====================================Function子类===================================
    
#负数
class Neg(Function):
    def forward(self,x):
        return -x
    def backward(self,gy):
        return -gy
def neg(x):
    return Neg()(x)

class Add(Function):
    """
    相加 参数和返回值都是列表
    """
    def forward(self,x0,x1):
        self.x0_shape, self.x1_shape = x0.shape, x1.shape
        y = x0 + x1
        return y
    def backward(self,gy):
        gx0, gx1 = gy, gy
        if self.x0_shape != self.x1_shape:
            # 若shape不同,则正向相加时进行了广播,反向传播时使用sum_to
            gx0 = dezero.functions.sum_to(gx0,self.x0_shape)
            gx1 = dezero.functions.sum_to(gx1, self.x1_shape)
        return gx0,gx1
def add(x0,x1):
    x1 = as_array(x1, dezero.cuda.get_array_module(x0.data)) #__add__可输入int float等类型
    return Add()(x0,x1)
    
class Mul(Function):
    def forward(self,x0,x1):
        y = x0 * x1
        return y
    def backward(self,gy):
        x0, x1 = self.inputs
        return gy*x1, gy*x0
def mul(x0,x1):
    x1 = as_array(x1, dezero.cuda.get_array_module(x0.data))
    return Mul()(x0,x1)
    
class Sub(Function):
    def forward(self,x0,x1):
        y = x0 - x1
        return y
    def backward(self,gy):
        return gy, -gy
def sub(x0,x1):
    x1 = as_array(x1, dezero.cuda.get_array_module(x0.data))
    return Sub()(x0,x1)
def rsub(x0,x1):
    x1 = as_array(x1, dezero.cuda.get_array_module(x0.data))
    return Sub()(x1,x0) #交换位置
    
class Div(Function):
    def forward(self,x0,x1):
        y = x0 / x1
        return y
    def backward(self,gy):
        x0, x1 = self.inputs
        gx0 = gy / x1
        gx1 = gy * (-x0 / x1**2)
        return gx0,gx1
def div(x0,x1):
    x1 = as_array(x1, dezero.cuda.get_array_module(x0.data))
    return Div()(x0,x1)
def rdiv(x0,x1):
    x1 = as_array(x1, dezero.cuda.get_array_module(x0.data))
    return Div()(x1,x0)
    
class Pow(Function):
    def __init__(self,c):
        self.c = c
    def forward(self,x):
        y = x ** self.c
        return y
    def backward(self, gy):
        x, = self.inputs #注意这边返回的是元组
        c = self.c
        gx = gy * c * x ** (c-1)
        return gx
def pow(x,c):
    return Pow(c)(x)
    

def setup_variable():
    """运算符重载"""
    Variable.__mul__ = mul
    Variable.__rmul__ = mul
    Variable.__add__ = add
    Variable.__radd__ = add
    Variable.__neg__ = neg
    Variable.__sub__ = sub
    Variable.__rsub__ = rsub
    Variable.__truediv__ = div
    Variable.__rtruediv__ = rdiv
    Variable.__pow__ = pow

    # Variable.matmul = dezero.functions.matmul
    # Variable.dot = dezero.functions.matmul
    # Variable.max = dezero.functions.max
    # Variable.min = dezero.functions.min

#================================
# 神经网络层
#================================

class Parameter(Variable):
    """神经网络参数 参数同Variable(data,name)"""
    pass









